use crate::cryptography::{dk, hmac_md5, hmac_sha1, md5, AesSizes};

pub fn checksum_hmac_md5(
    key: &[u8],
    key_usage: i32,
    plaintext: &[u8],
) -> Vec<u8> {
    let mut keyword = "signaturekey".to_string().into_bytes();
    keyword.push(0);

    let ksign = hmac_md5(key, &keyword);
    let mut bs = key_usage.to_le_bytes().to_vec();
    bs.append(&mut plaintext.to_vec());
    let tmp = md5(&bs);

    return hmac_md5(&ksign, &tmp);
}

pub fn checksum_sha_aes(
    key: &[u8],
    key_usage: i32,
    plaintext: &[u8],
    aes_sizes: &AesSizes,
) -> Vec<u8> {
    let mut key_usage_bytes = key_usage.to_be_bytes().to_vec();
    key_usage_bytes.push(0x99);

    let kc = dk(key, &key_usage_bytes, aes_sizes);

    let mac = hmac_sha1(&kc, plaintext);

    return mac[..12].to_vec();
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_checksum_hmac_md5() {
        let s4u_byte_array = &[
            0x01, 0x00, 0x00, 0x00, 0x61, 0x64, 0x6D, 0x69, 0x6E, 0x69, 0x73,
            0x74, 0x72, 0x61, 0x64, 0x6F, 0x72, 0x6B, 0x69, 0x6E, 0x67, 0x64,
            0x6F, 0x6D, 0x2E, 0x68, 0x65, 0x61, 0x72, 0x74, 0x73, 0x4B, 0x65,
            0x72, 0x62, 0x65, 0x72, 0x6F, 0x73,
        ];

        let session_key = &[
            0x72, 0xC3, 0x90, 0xC6, 0x5D, 0x71, 0x89, 0xAB, 0x9A, 0x39, 0xC3,
            0xFB, 0xFC, 0xBA, 0x41, 0xB8, 0x5A, 0x5F, 0x72, 0x6E, 0xD6, 0x4C,
            0x0C, 0x8A, 0x5C, 0xAB, 0x5B, 0xEB, 0x64, 0x8C, 0xB8, 0x74,
        ];

        let result = vec![
            0x85, 0xBF, 0xBD, 0x75, 0xEE, 0x4A, 0x87, 0x92, 0x9C, 0x28, 0xBC,
            0x07, 0x24, 0x95, 0x2F, 0x8D,
        ];

        assert_eq!(result, checksum_hmac_md5(session_key, 17, s4u_byte_array))
    }

    #[test]
    fn test_checksum_sha_aes256() {
        let key = &[
            0x25, 0xB2, 0x07, 0x6C, 0xDA, 0x3B, 0xFD, 0x62, 0x09, 0x16, 0x1A,
            0x6C, 0x78, 0xA6, 0x9C, 0x1C, 0x25, 0xB2, 0x07, 0x6C, 0xDA, 0x3B,
            0xFD, 0x62, 0x09, 0x16, 0x1A, 0x6C, 0x78, 0xA6, 0x9C, 0x1C,
        ];

        let blob: &'static [u8] = &[
            0xFF, 0xFF, 0xFF, 0x7F, 0x80, 0x23, 0xB8, 0xA3, 0x10, 0x31, 0xD6,
            0x01, 0x00, 0x00, 0x00, 0x00, 0x45,
        ];

        assert_eq!(
            vec![
                0xA8, 0x4D, 0x32, 0x1D, 0x23, 0x24, 0xFF, 0xA0, 0x62, 0x44,
                0xB7, 0x87
            ],
            checksum_sha_aes(key, 17, blob, &AesSizes::Aes256)
        );
    }

    #[test]
    fn test_checksum_sha_aes128() {
        let key = &[
            0x25, 0xB2, 0x07, 0x6C, 0xDA, 0x3B, 0xFD, 0x62, 0x09, 0x16, 0x1A,
            0x6C, 0x78, 0xA6, 0x9C, 0x1C,
        ];

        let blob: &'static [u8] = &[
            0xFF, 0xFF, 0xFF, 0x7F, 0x80, 0x23, 0xB8, 0xA3, 0x10, 0x31, 0xD6,
            0x01, 0x00, 0x00, 0x00, 0x00, 0x45,
        ];

        assert_eq!(
            vec![
                0xF0, 0x5F, 0xEF, 0x82, 0x1E, 0xCD, 0x6D, 0x62, 0xAD, 0x3A,
                0x38, 0x94
            ],
            checksum_sha_aes(key, 17, blob, &AesSizes::Aes128)
        );
    }
}
